{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "T4",
      "cell_execution_strategy": "setup",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/eyaler/avatars4all/blob/master/facevidcrop.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#FaceVidCrop\n",
        "\n",
        "### Notebook by [Eyal Gruss](https://eyalgruss.com), [@eyaler](twitter.com/eyaler)\n",
        "\n",
        "Crop a video to center on a face + optionally upscale/restore with [GFPGAN](https://github.com/TencentARC/GFPGAN)\n",
        "\n",
        "More notebooks: [github.com/eyaler/avatars4all](https://github.com/eyaler/avatars4all)\n",
        "\n",
        "Shortcut here: [tfi.la/face](https://tfi.la/face)\n",
        "\n",
        "Something not working? Open an [issue](https://github.com/eyaler/avatars4all/issues)\n",
        "\n",
        "If you find my work useful please consider supporting me via [GitHub Sponsors](https://github.com/sponsors/eyaler) or [PayPal](https://www.paypal.com/donate/?hosted_button_id=LNJ6F3FR79ARE)"
      ],
      "metadata": {
        "id": "fIHL9q7CJ4tR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Setup\n",
        "\n",
        "import locale\n",
        "locale.getpreferredencoding = lambda: 'UTF-8'\n",
        "\n",
        "!pip install git+https://github.com/ytdl-org/youtube-dl\n",
        "!pip install git+https://github.com/1adrianb/face-alignment\n",
        "\n",
        "%cd /content\n",
        "!git clone --depth=1 https://github.com/TencentARC/GFPGAN\n",
        "%cd /content/GFPGAN\n",
        "!pip install basicsr\n",
        "!pip install facexlib\n",
        "!pip install -r requirements.txt\n",
        "!python setup.py develop\n",
        "!pip install realesrgan\n",
        "%cd /content"
      ],
      "metadata": {
        "cellView": "form",
        "id": "NrVPvyVku1BJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Optionally mount Google Drive (MARK CHECKBOX) { run: \"auto\" }\n",
        "mount_google_drive = False #@param {type:\"boolean\"}\n",
        "if mount_google_drive:\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/drive')\n",
        "  print('path is /content/drive/MyDrive')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "ods_VYgxFFaU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ANINpF9JKuk8",
        "cellView": "form"
      },
      "source": [
        "#@title Crop video to face\n",
        "\n",
        "video_url = 'https://www.youtube.com/watch?v=vfIHsurwTMo' #@param {type: 'string'}\n",
        "#@markdown (leave empty to upload file - A BUTTON WILL APPEAR BELOW, or link to youtube / vimeo / video url / path to video on mounted drive [drive/MyDrive/...] / path to video on colab)\n",
        "i_just_uploaded_a_file_and_i_want_to_reuse_that_instead_of_uploading_a_new_one = False  #@param {type: 'boolean'}\n",
        "face_order = 'left-to-right then top-to-bottom' #@param ['left-to-right then top-to-bottom', 'right-to-left then top-to-bottom', 'top-to-bottom then left-to-right', 'bottom-to-top then left-to-right']\n",
        "#@markdown (one of these orderings may be more stable for your video, as order is recalculated per frame)\n",
        "face_num = 0 #@param {type: 'integer'}\n",
        "#@markdown (use face_num = 0 to select the largest face, face_num >= 1 to select a face from all the faces ordered as above)\n",
        "discard_faces_with_longest_dimension_pixels_below = 0 #@param {type: 'integer'}\n",
        "min_conf = 0.9 #@param {type: 'number'}\n",
        "#@markdown (lower min_conf if faces are not detected, raise min_conf if there are false detections that you cannot get rid of by using discard_faces_with_longest_dimension_pixels_below)\n",
        "temporal_smoothing = 0.9 #@param {type: 'number'}\n",
        "#@markdown (use temporal_smoothing = 0 to track head without smoothing, temporal_smoothing = 1 will not track head movement)\n",
        "top_extend_frac = 0.33 #@param {type: 'number'}\n",
        "bottom_extend_frac = 0.33 #@param {type: 'number'}\n",
        "#@markdown (leave empty to extend to the top/bottom of the frame)\n",
        "aspect_ratio = 1.333 #@param {type: 'number'}\n",
        "#@markdown (use aspect_ratio = 0 to grab from input video)\n",
        "start_seconds = 0 #@param {type: 'number'}\n",
        "duration_seconds = 0 #@param {type: 'number'}\n",
        "#@markdown (use duration_seconds = 0 for unrestricted duration)\n",
        "GFPGAN_model = '1.4' #@param ['1', '1.2', '1.3', '1.4', 'RestoreFormer']\n",
        "GFPGAN_factor = 0 #@param {type: 'integer'}\n",
        "#@markdown (use GFPGAN_factor = 0 to skip GFPGAN, GFPGAN_factor = 1 for restore without upscaling, GFPGAN_factor >= 2 for upscaling)\n",
        "max_width = 0 #@param {type: 'integer'}\n",
        "#@markdown (use max_width = 0 for unrestricted width)\n",
        "max_height = 0 #@param {type: 'integer'}\n",
        "#@markdown (use max_height = 0 for unrestricted height)\n",
        "copy_audio = True  #@param {type: 'boolean'}\n",
        "temporal_mode = 'forward' #@param ['forward', 'reverse', 'forward + reverse']\n",
        "output_filepath = 'output.mp4' #@param {type: 'string'}\n",
        "#@markdown (you can specify a path to save on mounted drive [drive/MyDrive/...])\n",
        "\n",
        "start_seconds = start_seconds or 0\n",
        "duration_seconds = duration_seconds or 0\n",
        "max_width = max_width or 0\n",
        "max_height = max_height or 0\n",
        "\n",
        "from time import time\n",
        "start_time = time()\n",
        "\n",
        "%cd /content\n",
        "\n",
        "import os\n",
        "from google.colab import files\n",
        "output_filepath = os.path.normpath(output_filepath.replace('\\\\', '/'))\n",
        "output_path, output_filename = os.path.split(output_filepath)\n",
        "\n",
        "need_dl = True\n",
        "try:\n",
        "  if orig_video and video_url == save_url and (video_url or i_just_uploaded_a_file_and_i_want_to_reuse_that_instead_of_uploading_a_new_one):\n",
        "    need_dl = False\n",
        "except:\n",
        "  pass\n",
        "if need_dl:\n",
        "  if not video_url:\n",
        "    %cd /content/sample_data\n",
        "    try:\n",
        "      uploaded = files.upload()\n",
        "    except Exception:\n",
        "      %cd /content\n",
        "      raise\n",
        "    for fn in uploaded:\n",
        "      orig_video = os.path.abspath(fn)\n",
        "      break\n",
        "    else:\n",
        "      need_dl = False\n",
        "    %cd /content\n",
        "  elif os.path.isfile(os.path.normpath(video_url.replace('\\\\', '/'))):\n",
        "    orig_video = os.path.abspath(os.path.normpath(video_url.replace('\\\\', '/')))\n",
        "  elif os.path.isfile(os.path.normpath('/content/drive/MyDrive/' + video_url.replace('\\\\', '/'))):\n",
        "    orig_video = os.path.normpath('/content/drive/MyDrive/' + video_url.replace('\\\\', '/'))\n",
        "  else:\n",
        "    orig_video = '/content/orig_video.mp4'\n",
        "    !rm -f $orig_video\n",
        "    !youtube-dl --no-check-extensions --no-playlist -f \"bestvideo[ext=mp4][vcodec!*=av01]+bestaudio[ext=m4a]/mp4[vcodec!*=av01]/bestvideo[ext=mp4]+bestaudio[ext=m4a]/bestvideo[ext=mp4]+bestaudio/mp4\" \"$video_url\" --merge-output-format mp4 -o $orig_video\n",
        "    if not os.path.exists(orig_video):\n",
        "      orig_video = '/content/orig_video.' + video_url.rsplit('.', 1)[1]\n",
        "      !wget \"$video_url\" -O $orig_video\n",
        "      if not os.path.getsize(orig_video):\n",
        "        !rm -f $orig_video\n",
        "  if not os.path.exists(orig_video):\n",
        "    orig_video = None\n",
        "    raise FileNotFoundError\n",
        "  input_video = '/content/input_video' + os.path.splitext(orig_video)[-1]\n",
        "\n",
        "need_fix_duration = True\n",
        "try:\n",
        "  if not need_dl and start_seconds == save_start_seconds and duration_seconds == save_duration_seconds:\n",
        "    need_fix_duration = False\n",
        "except:\n",
        "  pass\n",
        "if need_fix_duration:\n",
        "  if start_seconds or duration_seconds:\n",
        "    !ffmpeg -y -ss $start_seconds -t $duration_seconds -i \"$orig_video\" -f mp4 $input_video\n",
        "  else:\n",
        "    !cp \"$orig_video\" $input_video\n",
        "\n",
        "save_url = video_url\n",
        "save_start_seconds = start_seconds\n",
        "save_duration_seconds = duration_seconds\n",
        "\n",
        "import cv2\n",
        "import face_alignment\n",
        "face_detector_kwargs = {}\n",
        "if min_conf < 0.5:\n",
        "  face_detector_kwargs = {'filter_threshold': min_conf}\n",
        "fa = face_alignment.FaceAlignment(landmarks_type=1, face_detector_kwargs=face_detector_kwargs)\n",
        "import imageio.v3 as iio\n",
        "from IPython.display import display, Image, Video, clear_output\n",
        "\n",
        "!rm -f /content/numbers.png\n",
        "!rm -rf /content/in_frames\n",
        "!mkdir -p /content/in_frames\n",
        "!rm -rf /content/out_frames\n",
        "!mkdir -p /content/out_frames\n",
        "\n",
        "faces = []\n",
        "good_faces = []\n",
        "first_faces = ''\n",
        "max_face_len_unfiltered = 0\n",
        "max_face_len_below = 0\n",
        "min_face_len_above = None\n",
        "max_face_conf_unfiltered = 0\n",
        "max_face_conf_below = 0\n",
        "min_face_conf_above = 1\n",
        "max_filtered_faces_in_frame = 0\n",
        "frame_of_max_filtered_faces_in_frame = 0\n",
        "for frame, im in enumerate(iio.imiter(input_video)):\n",
        "  if min_face_len_above is None:\n",
        "    min_face_len_above = max(im.shape[:2])\n",
        "  *_, bboxes = fa.get_landmarks_from_image(im, return_bboxes=True)\n",
        "  if bboxes is None:\n",
        "    bboxes = []\n",
        "  if len(bboxes):\n",
        "    max_face_len_unfiltered = max(max_face_len_unfiltered, *[max(x1 - x0, y1 - y0) + 1 for x0, y0, x1, y1, _ in bboxes])\n",
        "    max_face_conf_unfiltered = max(max_face_conf_unfiltered, *[b[-1] for b in bboxes])\n",
        "\n",
        "  if discard_faces_with_longest_dimension_pixels_below:\n",
        "    for x0, y0, x1, y1, _ in bboxes:\n",
        "      m = max(x1 - x0, y1 - y0) + 1\n",
        "      if m < discard_faces_with_longest_dimension_pixels_below:\n",
        "        max_face_len_below = max(max_face_len_below, m)\n",
        "    bboxes = [b for b in bboxes if max(b[2] - b[0], b[3] - b[1]) + 1 >= discard_faces_with_longest_dimension_pixels_below]\n",
        "  if len(bboxes):\n",
        "    min_face_len_above = min(min_face_len_above, *[max(x1 - x0, y1 - y0) + 1 for x0, y0, x1, y1, _ in bboxes])\n",
        "\n",
        "  if min_conf:\n",
        "    for b in bboxes:\n",
        "      if b[-1] < min_conf:\n",
        "        max_face_conf_below = max(max_face_conf_below, b[-1])\n",
        "    bboxes = [b for b in bboxes if b[-1] >= min_conf]\n",
        "  if len(bboxes):\n",
        "    min_face_conf_above = min(min_face_conf_above, *[b[-1] for b in bboxes])\n",
        "  bboxes = sorted(bboxes, key=lambda b: (b[0], b[2], b[1], b[3]) if face_order == 'left-to-right then top-to-bottom' else (-b[2], -b[0], b[1], b[3]) if face_order == 'right-to-left then top-to-bottom' else (b[1], b[3], b[0], b[2]) if face_order == 'top-to-bottom then left-to-right' else (-b[3], -b[1], b[0], b[2]))\n",
        "  if len(bboxes) > max_filtered_faces_in_frame:\n",
        "    max_filtered_faces_in_frame = len(bboxes)\n",
        "    frame_of_max_filtered_faces_in_frame = frame\n",
        "  if len(bboxes) > 1 and not first_faces:\n",
        "    first_faces = f'{frame=}'\n",
        "    for i, (x0, y0, x1, y1, conf) in enumerate(bboxes):\n",
        "      cv2.putText(im, str(i + 1), (max(min(int(x1), im.shape[1] - 1) - 20, 0), max(min(int(y1), im.shape[0] - 1), 20)), 0, .7, (0, 255, 0), 2)\n",
        "      first_faces += f'; {i + 1}: w={x1 - x0 + 1 :.1f} h={y1 - y0 + 1 :.1f} {conf=:.2f}'\n",
        "    iio.imwrite('/content/numbers.png', im, compress_level=1)\n",
        "    clear_output()\n",
        "    display(Image('/content/numbers.png'))\n",
        "    print(first_faces)\n",
        "  if face_num:\n",
        "    bboxes = bboxes[face_num - 1 : face_num]\n",
        "  if bboxes:\n",
        "    bbox = sorted(bboxes, key=lambda b: (b[2]-b[0]+1) * (b[3]-b[1]+1))[-1]\n",
        "    faces.append(bbox[:-1])\n",
        "    good_faces.append(faces[-1])\n",
        "  else:\n",
        "    faces.append(None)\n",
        "\n",
        "print(f'orig_video={orig_video}')\n",
        "print(f'max_face_len_below_{discard_faces_with_longest_dimension_pixels_below}={max_face_len_below:.1f} min_face_len_above_{discard_faces_with_longest_dimension_pixels_below}={min_face_len_above:.1f} {max_face_len_unfiltered=:.1f}')\n",
        "print(f'max_face_conf_below_{min_conf}={max_face_conf_below:.2f} min_face_conf_above_{min_conf}={min_face_conf_above:.2f} {max_face_conf_unfiltered=:.2f}')\n",
        "assert_str = f'No faces found. Consider decreasing min_conf below {min(min_conf, max_face_conf_unfiltered):.2f}'\n",
        "if discard_faces_with_longest_dimension_pixels_below:\n",
        "  assert_str += f' or discard_faces_with_longest_dimension_pixels_below below {min(discard_faces_with_longest_dimension_pixels_below, max_face_len_unfiltered):.1f}'\n",
        "assert good_faces, assert_str\n",
        "median = sorted(good_faces, key=lambda b: (b[2]-b[0]+1) * (b[3]-b[1]+1))[len(good_faces) // 2]\n",
        "if top_extend_frac is None:\n",
        "  top = (median[3]+median[1]) / 2\n",
        "else:\n",
        "  top = (median[3]-median[1]+1) * (.5+top_extend_frac)\n",
        "if bottom_extend_frac is None:\n",
        "  bottom = im.shape[0] - (median[3]+median[1])/2\n",
        "else:\n",
        "  bottom = (median[3]-median[1]+1) * (.5+bottom_extend_frac)\n",
        "h = top + bottom\n",
        "if h > im.shape[0]:\n",
        "  top = top / h * im.shape[0]\n",
        "  h = im.shape[0]\n",
        "if aspect_ratio:\n",
        "  w = min(round(h * aspect_ratio), im.shape[1])\n",
        "else:\n",
        "  w = round(h * im.shape[1] / im.shape[0])\n",
        "\n",
        "prev_face = good_faces[0]\n",
        "for frame, (im, face) in enumerate(zip(iio.imiter(input_video), faces)):\n",
        "  x0, y0, x1, y1 = prev_face = prev_face if face is None else face\n",
        "\n",
        "  if frame:\n",
        "    x = x*temporal_smoothing + (x0+x1)/2*(1-temporal_smoothing)\n",
        "    y = y*temporal_smoothing + (y0+y1)/2*(1-temporal_smoothing)\n",
        "  else:\n",
        "    x = (x0+x1) / 2\n",
        "    y = (y0+y1) / 2\n",
        "\n",
        "  x0 = max(x - w/2, 0)\n",
        "  x1 = int(min(x0 + w, im.shape[1]) - 1)\n",
        "  x0 = int(x1 - w + 1)\n",
        "\n",
        "  y0 = max(y - top, 0)\n",
        "  y1 = int(min(y0 + h, im.shape[0]) - 1)\n",
        "  y0 = int(y1 - h + 1)\n",
        "\n",
        "  im = im[y0:y1, x0:x1]\n",
        "  iio.imwrite(f'/content/in_frames/frame_{frame:06d}.png', im, compress_level=1)\n",
        "\n",
        "reverse_filename = 'reverse_' + output_filename\n",
        "if GFPGAN_factor:\n",
        "  %cd /content/GFPGAN\n",
        "  !python inference_gfpgan.py -i /content/in_frames -o /content/out_frames -v $GFPGAN_model -s $GFPGAN_factor\n",
        "  %cd /content\n",
        "  im_folder = 'out_frames/restored_imgs'\n",
        "else:\n",
        "  im_folder = 'in_frames'\n",
        "\n",
        "fps = iio.immeta(input_video)['fps']\n",
        "if 'forward' in temporal_mode:\n",
        "  if copy_audio:\n",
        "    !ffmpeg -y -framerate $fps -thread_queue_size 0 -i /content/$im_folder/frame_%06d.png -i $input_video -c:v libx264 -c:a aac -map 0:v -map 1:a? -vf \"scale=min(iw\\,$max_width):min(ih\\,$max_height):force_original_aspect_ratio=decrease:force_divisible_by=2\" -pix_fmt yuv420p -profile:v baseline -movflags +faststart \"/content/$output_filename\"\n",
        "  else:\n",
        "    !ffmpeg -y -framerate $fps -thread_queue_size 0 -i /content/$im_folder/frame_%06d.png -c:v libx264 -vf \"scale=min(iw\\,$max_width):min(ih\\,$max_height):force_original_aspect_ratio=decrease:force_divisible_by=2\" -pix_fmt yuv420p -profile:v baseline -movflags +faststart \"/content/$output_filename\"\n",
        "\n",
        "if 'reverse' in temporal_mode:\n",
        "  if copy_audio:\n",
        "    !ffmpeg -y -framerate $fps -thread_queue_size 0 -i /content/$im_folder/frame_%06d.png -i $input_video -c:v libx264 -c:a aac -map 0:v -map 1:a? -vf \"scale=min(iw\\,$max_width):min(ih\\,$max_height):force_original_aspect_ratio=decrease:force_divisible_by=2,reverse\" -af areverse -pix_fmt yuv420p -profile:v baseline -movflags +faststart \"/content/$reverse_filename\"\n",
        "  else:\n",
        "    !ffmpeg -y -framerate $fps -thread_queue_size 0 -i /content/$im_folder/frame_%06d.png -c:v libx264 -vf \"scale=min(iw\\,$max_width):min(ih\\,$max_height):force_original_aspect_ratio=decrease:force_divisible_by=2,reverse\" -af areverse -pix_fmt yuv420p -profile:v baseline -movflags +faststart \"/content/$reverse_filename\"\n",
        "\n",
        "if temporal_mode == 'reverse':\n",
        "  !mv \"/content/$reverse_filename\" \"/content/$output_filename\"\n",
        "elif temporal_mode == 'forward + reverse':\n",
        "  !echo file \"/content/$output_filename\" > list.txt\n",
        "  !echo file \"/content/$reverse_filename\" >> list.txt\n",
        "  concat_filename = 'concat_' + output_filename\n",
        "  !ffmpeg -y -f concat -safe 0 -i list.txt -c copy \"/content/$concat_filename\"\n",
        "  !mv \"/content/$concat_filename\" \"/content/$output_filename\"\n",
        "\n",
        "clear_output()\n",
        "if first_faces:\n",
        "  display(Image('/content/numbers.png'))\n",
        "  print(first_faces)\n",
        "if output_path:\n",
        "  !mkdir -p \"$output_path\"\n",
        "  !cp \"/content/$output_filename\" \"$output_filepath\"\n",
        "meta = iio.immeta('/content/' + output_filename)\n",
        "print(f'took {(time()-start_time) / 60 :.1f} min. orig_video={orig_video}')\n",
        "print(f'output_video=/content/{output_filename} w={meta[\"size\"][0]} h={meta[\"size\"][1]} aspect_ratio={meta[\"size\"][0] / meta[\"size\"][1] :.3f} t={meta[\"duration\"]} fps={meta[\"fps\"]}')\n",
        "print(f'max_face_len_below_{discard_faces_with_longest_dimension_pixels_below}={max_face_len_below:.1f} min_face_len_above_{discard_faces_with_longest_dimension_pixels_below}={min_face_len_above:.1f} {max_face_len_unfiltered=:.1f}')\n",
        "print(f'max_face_conf_below_{min_conf}={max_face_conf_below:.2f} min_face_conf_above_{min_conf}={min_face_conf_above:.2f} {max_face_conf_unfiltered=:.2f}')\n",
        "print(f'{max_filtered_faces_in_frame=} {frame_of_max_filtered_faces_in_frame=}')\n",
        "print('if video does not show below, you can still download it!')\n",
        "display(Video('/content/' + output_filename, embed=True, html_attributes=\"autoplay controls loop\"))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lH4oJcBOKZcf",
        "cellView": "form"
      },
      "source": [
        "#@title Download\n",
        "\n",
        "from google.colab import files\n",
        "files.download('/content/' + output_filename)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}
